import asyncio
import json
from contextlib import AsyncExitStack
from typing import Optional

from mcp import StdioServerParameters, ClientSession
from mcp.client.stdio import stdio_client
from openai import OpenAI


class MCPClient:

    server_path= "server.py"

    def __init__(self):
        #创建线程管理栈
        self.exit_stack = AsyncExitStack()
        self.session: Optional[ClientSession] = None
        self.deepseek = OpenAI(
            api_key="sk-15af4e21f828460683b16ce9e78b2346",
            base_url="https://api.deepseek.com"
        )

    """创建连接服务端"""
    async def connect_to_server(self, server_path: str):
        # 一、创建StdioServerParameter参数信息
        server_parameters = StdioServerParameters(
            command="python",
            args=[server_path],
            env=None
        )
        # 二、 创建stdioClient
        client = stdio_client(server=server_parameters)
        stdio_transport = await self.exit_stack.enter_async_context(client)
        read_stream, write_stream = stdio_transport

        # 三、创建ClientSession
        client_session = ClientSession(read_stream, write_stream)
        self.session = await self.exit_stack.enter_async_context(client_session)
        # 四、初始化session
        await self.session.initialize()
        #五、列出所有可用的工具
        # response = await self.session.list_tools()
        # tools = response.tools
        # print("\nConnected to server with tools:",tools)

    async def execute(self,query:str):
        # 一、通过session列表所有的工具 及组装function_calling
        response = await self.session.list_tools()
        tools = response.tools
        print("\nConnected to server with tools:",tools)
        # 组装function calling
        tools = [
            {"type":"function",
             "function":{
                 "name":tool.name,
                 "description":tool.description,
                 "parameters":tool.inputSchema
                }
             } for tool in tools
        ]
        # 二、 大模型调用参数：组装messages
        messages = [
            {
                "role": "user",
                "content": query
            }
        ]
        # 调用大模型
        deepseek_response = self.deepseek.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
            tools=tools,
        )
        # 获取大模型的决策结果
        print("==== deepseek 决策结果：",deepseek_response)
        choice_result = deepseek_response.choices[0]
        print("最终选择的工具：",choice_result.message.tool_calls[0].function.name)

        '''  ******************* 3.4 function calling 过程 *******************'''
        if choice_result.finish_reason == "tool_calls":
            # 根据大模型选择的工具，进行调用执行
            messages.append(choice_result.message.model_dump())
            # 调用工具链
            tool_call = choice_result.message.tool_calls[0]
            function_name = tool_call.function.name
            arguments = json.loads(tool_call.function.arguments)

            '''  调用function calling 工具链'''
            tool_result = await self.session.call_tool(
                name = function_name,
                arguments=arguments
            )

            print("==== 工具调用结果元数据：",tool_result)
            print("====  工具计算结果：",tool_result.content[0].text)







    #关闭连接
    async def cleanup(self):
        await self.exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.connect_to_server(client.server_path)
        await client.execute("计算1+2")
    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())